#include "TcpServer.h"
#include <ws2tcpip.h>
#include <iostream>

TcpServer::TcpServer() : server_fd(INVALID_SOCKET), client_fd(INVALID_SOCKET) {
    WSADATA wsaData;
    if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0) {
        error_msg = "WSAStartup failed";
    }
}

TcpServer::~TcpServer() {
    Close();
}

bool TcpServer::Initialize(const std::string& ip, int port) {
    server_fd = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
    if (server_fd == INVALID_SOCKET) {
        error_msg = "Socket creation failed";
        return false;
    }

    sockaddr_in server_addr{};
    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(port);
    inet_pton(AF_INET, ip.c_str(), &server_addr.sin_addr);

    if (bind(server_fd, (sockaddr*)&server_addr, sizeof(server_addr)) == SOCKET_ERROR) {
        error_msg = "Bind failed";
        Close();
        return false;
    }

    if (listen(server_fd, 3) == SOCKET_ERROR) {
        error_msg = "Listen failed";
        Close();
        return false;
    }

    return true;
}

bool TcpServer::Send(const std::string& data, const std::string&, int) {
    if (client_fd == INVALID_SOCKET) {
        error_msg = "No client connected";
        return false;
    }
    int sent = send(client_fd, data.c_str(), data.length(), 0);
    if (sent == SOCKET_ERROR) {
        error_msg = "Send failed";
        return false;
    }
    return true;
}

bool TcpServer::Receive(std::function<void(const std::string&, const std::string&, int)> callback) {
    sockaddr_in client_addr{};
    int addr_len = sizeof(client_addr);
    client_fd = accept(server_fd, (sockaddr*)&client_addr, &addr_len);
    if (client_fd == INVALID_SOCKET) {
        error_msg = "Accept failed";
        return false;
    }

    char buffer[1024] = { 0 };
    int received = recv(client_fd, buffer, sizeof(buffer), 0);
    if (received > 0) {
        char client_ip[INET_ADDRSTRLEN];
        inet_ntop(AF_INET, &client_addr.sin_addr, client_ip, INET_ADDRSTRLEN);
        callback(std::string(buffer, received), client_ip, ntohs(client_addr.sin_port));
    }
    else {
        error_msg = "Receive failed";
        closesocket(client_fd);
        client_fd = INVALID_SOCKET;
        return false;
    }
    return true;
}

void TcpServer::Close() {
    if (client_fd != INVALID_SOCKET) {
        closesocket(client_fd);
        client_fd = INVALID_SOCKET;
    }
    if (server_fd != INVALID_SOCKET) {
        closesocket(server_fd);
        server_fd = INVALID_SOCKET;
    }
    WSACleanup();
}

std::string TcpServer::GetError() const {
    return error_msg;
}